{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2025-10-09T03:09:38.486469Z",
     "start_time": "2025-10-09T03:09:31.992705Z"
    }
   },
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "\n",
    "from d2l_learn.utils import masked_softmax\n",
    "\n",
    "\n",
    "# 加性注意力函数\n",
    "class AdditiveAttention(nn.Module):\n",
    "    def __init__(self, query_size, key_size, num_hiddens, dropout=0.1):\n",
    "        super(AdditiveAttention, self).__init__()\n",
    "        self.W_q = nn.Linear(query_size, num_hiddens, bias=False)\n",
    "        self.W_k = nn.Linear(key_size, num_hiddens, bias=False)\n",
    "        self.w_v = nn.Linear(num_hiddens, 1, bias=False)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "    def forward(self, q, k, v, valid_lens):\n",
    "        # q的形状为(batch_size, 查询的个数, query_size)\n",
    "        # k的形状为(batch_size, 键-值对的个数, key_size)\n",
    "        # v的形状为(batch_size, 键-值对的个数, value_size)\n",
    "        q = self.W_q(q) # q 的形状为(batch_size, 查询的个数, hidden_size)\n",
    "        k = self.W_k(k) # k 的形状为(batch_size, 键-值对的个数, hidden_size)\n",
    "        features = q.unsqueeze(2) + k.unsqueeze(1)\n",
    "        features = torch.tanh(features)\n",
    "        scores = self.w_v(features).squeeze(-1)\n",
    "        attention_weights = masked_softmax(scores, valid_lens)\n",
    "        return torch.bmm(self.dropout(attention_weights), v)"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-09T03:12:28.184746Z",
     "start_time": "2025-10-09T03:12:28.176470Z"
    }
   },
   "cell_type": "code",
   "source": [
    "queries, keys = torch.normal(0, 1, (2, 1, 20)), torch.ones((2, 10, 2))\n",
    "values = torch.arange(0, 40, dtype=torch.float32).reshape((1, 10, 4)).repeat((2, 1, 1))\n",
    "valid_lens = torch.tensor([2, 6], dtype=torch.long)\n",
    "attention = AdditiveAttention(queries.size(-1), keys.size(-1), 8)\n",
    "print(attention)\n",
    "attention.eval()\n",
    "print(attention(queries, keys, values, valid_lens))"
   ],
   "id": "1269c3a65275962d",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AdditiveAttention(\n",
      "  (W_q): Linear(in_features=20, out_features=8, bias=False)\n",
      "  (W_k): Linear(in_features=2, out_features=8, bias=False)\n",
      "  (w_v): Linear(in_features=8, out_features=1, bias=False)\n",
      "  (dropout): Dropout(p=0.1, inplace=False)\n",
      ")\n",
      "tensor([[[ 2.0000,  3.0000,  4.0000,  5.0000]],\n",
      "\n",
      "        [[10.0000, 11.0000, 12.0000, 13.0000]]], grad_fn=<BmmBackward0>)\n"
     ]
    }
   ],
   "execution_count": 7
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
